import json
from datetime import datetime, timezone
from unittest.mock import MagicMock, patch

import pytest

from openhands.core.config import AgentConfig, LLMConfig
from openhands.events.event import Event, EventSource
from openhands.events.stream import EventStream
from openhands.memory.long_term_memory import LongTermMemory
from openhands.storage.files import FileStore


@pytest.fixture
def mock_llm_config() -> LLMConfig:
    config = MagicMock(spec=LLMConfig)
    config.embedding_model = 'test_embedding_model'
    config.api_key = 'test_api_key'
    config.api_version = 'v1'
    return config


@pytest.fixture
def mock_agent_config() -> AgentConfig:
    config = AgentConfig(
        memory_enabled=True,
        memory_max_threads=4,
        llm_config='test_llm_config',
    )
    return config


@pytest.fixture
def mock_file_store() -> FileStore:
    store = MagicMock(spec=FileStore)
    store.sid = 'test_session'
    return store


@pytest.fixture
def mock_event_stream(mock_file_store: FileStore) -> EventStream:
    with patch('openhands.events.stream.EventStream') as MockEventStream:
        instance = MockEventStream.return_value
        instance.sid = 'test_session'
        instance.get_events = MagicMock()
        return instance


@pytest.fixture
def long_term_memory(
    mock_llm_config: LLMConfig,
    mock_agent_config: AgentConfig,
    mock_event_stream: EventStream,
) -> LongTermMemory:
    mod = LongTermMemory.__module__
    with patch(f'{mod}.chromadb.PersistentClient') as mock_chroma_client:
        mock_collection = MagicMock()
        mock_chroma_client.return_value.get_or_create_collection.return_value = (
            mock_collection
        )
        with (
            patch(f'{mod}.ChromaVectorStore', MagicMock()),
            patch(f'{mod}.EmbeddingsLoader', MagicMock()),
            patch(f'{mod}.VectorStoreIndex', MagicMock()),
        ):
            memory = LongTermMemory(
                llm_config=mock_llm_config,
                agent_config=mock_agent_config,
                event_stream=mock_event_stream,
            )
            memory.collection = mock_collection
            return memory


def _create_action_event(action: str) -> Event:
    """Helper function to create an action event."""
    event = Event()
    event._id = -1
    event._timestamp = datetime.now(timezone.utc).isoformat()
    event._source = EventSource.AGENT
    event.action = action
    return event


def _create_observation_event(observation: str) -> Event:
    """Helper function to create an observation event."""
    event = Event()
    event._id = -1
    event._timestamp = datetime.now(timezone.utc).isoformat()
    event._source = EventSource.ENVIRONMENT
    event.observation = observation
    return event


def test_add_event_with_action(long_term_memory: LongTermMemory):
    event = _create_action_event('test_action')
    long_term_memory._add_document = MagicMock()
    long_term_memory.add_event(event)
    assert long_term_memory.thought_idx == 1
    long_term_memory._add_document.assert_called_once()
    _, kwargs = long_term_memory._add_document.call_args
    assert kwargs['document'].extra_info['type'] == 'action'
    assert kwargs['document'].extra_info['id'] == 'test_action'


def test_add_event_with_observation(long_term_memory: LongTermMemory):
    event = _create_observation_event('test_observation')
    long_term_memory._add_document = MagicMock()
    long_term_memory.add_event(event)
    assert long_term_memory.thought_idx == 1
    long_term_memory._add_document.assert_called_once()
    _, kwargs = long_term_memory._add_document.call_args
    assert kwargs['document'].extra_info['type'] == 'observation'
    assert kwargs['document'].extra_info['id'] == 'test_observation'


def test_add_event_with_missing_keys(long_term_memory: LongTermMemory):
    # Creating an event with additional unexpected attributes
    event = Event()
    event._id = -1
    event._timestamp = datetime.now(timezone.utc).isoformat()
    event._source = EventSource.AGENT
    event.action = 'test_action'
    event.unexpected_key = 'value'

    long_term_memory._add_document = MagicMock()
    long_term_memory.add_event(event)
    assert long_term_memory.thought_idx == 1
    long_term_memory._add_document.assert_called_once()
    _, kwargs = long_term_memory._add_document.call_args
    assert kwargs['document'].extra_info['type'] == 'action'
    assert kwargs['document'].extra_info['id'] == 'test_action'


def test_events_to_docs_no_events(
    long_term_memory: LongTermMemory, mock_event_stream: EventStream
):
    mock_event_stream.get_events.side_effect = FileNotFoundError

    # convert events to documents
    documents = long_term_memory._events_to_docs()

    # since get_events raises, documents should be empty
    assert len(documents) == 0

    # thought_idx remains unchanged
    assert long_term_memory.thought_idx == 0


def test_load_events_into_index_with_invalid_json(
    long_term_memory: LongTermMemory, mock_event_stream: EventStream
):
    """Test loading events with malformed event data."""
    # Simulate an event that causes event_to_memory to raise a JSONDecodeError
    with patch(
        'openhands.memory.long_term_memory.event_to_memory',
        side_effect=json.JSONDecodeError('Expecting value', '', 0),
    ):
        event = _create_action_event('invalid_action')
        mock_event_stream.get_events.return_value = [event]

        # convert events to documents
        documents = long_term_memory._events_to_docs()

        # since event_to_memory raises, documents should be empty
        assert len(documents) == 0

    # thought_idx remains unchanged
    assert long_term_memory.thought_idx == 0


def test_embeddings_inserted_into_chroma(long_term_memory: LongTermMemory):
    event = _create_action_event('test_action')
    long_term_memory._add_document = MagicMock()
    long_term_memory.add_event(event)
    long_term_memory._add_document.assert_called()
    _, kwargs = long_term_memory._add_document.call_args
    assert 'document' in kwargs
    assert (
        kwargs['document'].text
        == '{"source": "agent", "action": "test_action", "args": {}}'
    )


def test_search_returns_correct_results(long_term_memory: LongTermMemory):
    mock_retriever = MagicMock()
    mock_retriever.retrieve.return_value = [
        MagicMock(get_text=MagicMock(return_value='result1')),
        MagicMock(get_text=MagicMock(return_value='result2')),
    ]
    with patch(
        'openhands.memory.long_term_memory.VectorIndexRetriever',
        return_value=mock_retriever,
    ):
        results = long_term_memory.search(query='test query', k=2)
        assert results == ['result1', 'result2']
        mock_retriever.retrieve.assert_called_once_with('test query')


def test_search_with_no_results(long_term_memory: LongTermMemory):
    mock_retriever = MagicMock()
    mock_retriever.retrieve.return_value = []
    with patch(
        'openhands.memory.long_term_memory.VectorIndexRetriever',
        return_value=mock_retriever,
    ):
        results = long_term_memory.search(query='no results', k=5)
        assert results == []
        mock_retriever.retrieve.assert_called_once_with('no results')


def test_add_event_increment_thought_idx(long_term_memory: LongTermMemory):
    event1 = _create_action_event('action1')
    event2 = _create_observation_event('observation1')
    long_term_memory.add_event(event1)
    long_term_memory.add_event(event2)
    assert long_term_memory.thought_idx == 2


def test_load_events_batch_insert(
    long_term_memory: LongTermMemory, mock_event_stream: EventStream
):
    event1 = _create_action_event('action1')
    event2 = _create_observation_event('observation1')
    event3 = _create_action_event('action2')
    mock_event_stream.get_events.return_value = [event1, event2, event3]

    # Mock insert_batch_docs
    with patch('openhands.utils.embeddings.insert_batch_docs') as mock_run_docs:
        # convert events to documents
        documents = long_term_memory._events_to_docs()

        # Mock the insert_batch_docs to simulate document insertion
        mock_run_docs.return_value = []

        # Call insert_batch_docs with the documents
        mock_run_docs(
            index=long_term_memory.index,
            documents=documents,
            num_workers=long_term_memory.memory_max_threads,
        )

        # Assert that insert_batch_docs was called with the correct arguments
        mock_run_docs.assert_called_once_with(
            index=long_term_memory.index,
            documents=documents,
            num_workers=long_term_memory.memory_max_threads,
        )

    # Check if thought_idx was incremented correctly
    assert long_term_memory.thought_idx == 3
